package com.inspur.edp.ai.flow.core.service;

import com.inspur.edp.ai.flow.api.FlowInvokeService;
import com.inspur.edp.ai.flow.api.FlowMetadataService;
import com.inspur.edp.ai.flow.api.FlowResponse;
import com.inspur.edp.ai.flow.api.ServiceContainer;
import com.inspur.edp.ai.flow.core.engine.*;
import com.inspur.edp.ai.flow.core.engine.databus.DefaultDataBus;
import com.inspur.edp.ai.flow.engine.*;
import com.inspur.edp.ai.flow.exception.FlowException;
import com.inspur.edp.ai.flow.metadata.AbsNode;
import com.inspur.edp.ai.flow.metadata.Parameter;
import io.iec.edp.caf.commons.utils.StringUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;

/**
 * @author lizhaorui
 * @date 2025/10/11
 * @description
 */

public class FlowInvokeServiceImpl implements FlowInvokeService {

    private final StatefulSupporter statefulSupporter;

    private final EventBus eventBus;

    private final FlowMetadataService metadataService;

    private FlowRTObjectGetter rtObjectGetter;

    public FlowInvokeServiceImpl() {
        this.statefulSupporter = ServiceContainer.getService(StatefulSupporter.class);
        this.eventBus = ServiceContainer.getService(EventBus.class);
        this.metadataService = ServiceContainer.getService(FlowMetadataService.class);
        this.rtObjectGetter = flowId -> {
            FlowRTObject memRuleRTObject = FlowRTCacheManager.getInstance().getRTObject(flowId);
            String version = metadataService.getVersion(flowId);
            if (!version.equals(memRuleRTObject.getVersion())) {
                FlowRTCacheManager.getInstance().clear(flowId);
                memRuleRTObject = FlowRTCacheManager.getInstance().getRTObject(flowId);
            }
            return memRuleRTObject;
        };
    }

    public Map<String, Object> invoke(String flowId, List<Object> args) {
        String flowInstanceId = UUID.randomUUID().toString();
        EngineContext context = this.buildEngineContext(flowId, flowInstanceId);
        context.setDataBus(new DefaultDataBus());
        AbsNode startNode = context.getFlowContext().getFlow().getStartNode();
        FlowExecuteEngine engine = new FlowExecuteEngine();
        engine.executeStartNode(context, startNode, args);
        return context.getReturnValues();
    }

    public Map<String, Object> invoke(String flowId, Map<String, Object> argMap) {
        String flowInstanceId = UUID.randomUUID().toString();
        EngineContext context = this.buildEngineContext(flowId, flowInstanceId);
        context.setDataBus(new DefaultDataBus());
        AbsNode startNode = context.getFlowContext().getFlow().getStartNode();
        List<Object> args = this.fromMapToList(argMap, startNode);
        FlowExecuteEngine engine = new FlowExecuteEngine();
        engine.executeStartNode(context, startNode, args);
        return context.getReturnValues();
    }

    public <T> T invokeWithOneReturnValue(String flowId, Map<String, Object> argMap, Class<T> cls) {
        Map<String, Object> returnMap = this.invoke(flowId, argMap);
        if (returnMap == null) {
            throw new FlowException("return Value parameter count  is 0 require is 1.");
        } else if (returnMap.size() != 1) {
            throw new FlowException("return Value parameter count  is " + returnMap.size() + " require is 1.");
        } else {
            for (Map.Entry<String, Object> entry : returnMap.entrySet()) {
                Object returnObject = entry.getValue();
                if (returnObject == null) {
                    return null;
                } else {
                    if (returnObject instanceof String) {
                        return (T) returnObject;
                    } else {
                        throw new FlowException("return value type is " + cls.getName() + " require is String type.");
                    }
                }
            }
            return null;
        }
    }

    public FlowResponse asyncInvoke(String flowId, String flowInstanceId, List<Object> args) {
        EngineContext engineContext = this.buildEngineContext(flowId, flowInstanceId);
        return this.internalAsyncInvoke(engineContext, args, false);
    }

    public FlowResponse stream(String flowId, String flowInstanceId, List<Object> args) {
        if (StringUtils.isEmpty(flowInstanceId)) {
            flowInstanceId = UUID.randomUUID().toString();
        }
        EngineContext engineContext = this.buildEngineContext(flowId, flowInstanceId);
        return this.internalAsyncInvoke(engineContext, args, true);
    }

    public FlowResponse stream(String flowId, String flowInstanceId, Map<String, Object> argMap) {
        EngineContext context = this.buildEngineContext(flowId, flowInstanceId);
        AbsNode startNode = context.getFlowContext().getFlow().getStartNode();
        List<Object> args = this.fromMapToList(argMap, startNode);
        return this.internalAsyncInvoke(context, args, true);
    }

    public List<String> getSuspendNodes(String flowId, String flowInstanceId) {
        DataBus dataBus = this.statefulSupporter.recovery(flowInstanceId);
        return (List<String>) dataBus.get(EngineConstant.SUSPEND_NODES_KEY_LABEL);
    }

    public FlowResponse internalAsyncInvoke(EngineContext context,
                                            List<Object> args,
                                            boolean streamOutput) {
        String flowInstanceId = context.getFlowContext().getFlowInstanceId();
        if (!StringUtils.isEmpty(flowInstanceId) && this.statefulSupporter.recovery(flowInstanceId) != null) {
            DataBus dataBus = this.statefulSupporter.recovery(flowInstanceId);
            context.setDataBus(dataBus);
            List<String> suspendNodes = (List<String>) dataBus.get(EngineConstant.SUSPEND_NODES_KEY_LABEL);
            if (suspendNodes == null || suspendNodes.size() == 0) {
                throw new FlowException("当前流程实例没有暂停的节点");
            }
            if (suspendNodes.size() > 1) {
                throw new FlowException("当前流程实例有多个暂停的节点");
            }
            String nodeId = suspendNodes.get(0);
            context.setDataBus(dataBus);
            if (streamOutput) {
                context.getFlowContext().setEventBus(eventBus);
                context.getFlowContext().setStreamOutput(true);
            }
            Callable<Map<String, Object>> task = () -> {
                FlowExecuteEngine engine = new FlowExecuteEngine();
                engine.executeMultiTurnNode(context, nodeId, args);
                return context.getReturnValues();
            };
            Future<Map<String, Object>> future = ThreadPoolService.getExecutor().submit(task);
            return new FlowResponse(flowInstanceId, this.eventBus, future);

        } else {
            context.setDataBus(new DefaultDataBus());
            AbsNode startNode = context.getFlowContext().getFlow().getStartNode();
            if (streamOutput) {
                context.getFlowContext().setEventBus(eventBus);
                context.getFlowContext().setStreamOutput(true);
            }
            Callable<Map<String, Object>> task = () -> {
                FlowExecuteEngine engine = new FlowExecuteEngine();
                engine.executeStartNode(context, startNode, args);
                return context.getReturnValues();
            };
            Future<Map<String, Object>> future = ThreadPoolService.getExecutor().submit(task);
            return new FlowResponse(flowInstanceId, this.eventBus, future);

        }
    }

    private List<Object> fromMapToList(Map<String, Object> argMap, AbsNode startNode) {
        List<Object> args = new ArrayList<>();
        if (argMap != null && startNode.getInputParams() != null) {
            for (Parameter inputParam : startNode.getInputParams()) {
                args.add(argMap.get(inputParam.getCode()));
            }
        }
        return args;
    }

    private EngineContext buildEngineContext(String flowId, String flowInstanceId) {
        FlowRTObject flowRTObject = this.rtObjectGetter.get(flowId);
        EngineContext engineContext = new EngineContext();
        // 构建flowContext
        Flow flow = new Flow();
        flow.setId(flowId);
        flow.setParentId(null);
        flow.setNodes(flowRTObject.getFlow().getNodes());
        flow.setEdges(flowRTObject.getFlow().getEdges());
        FlowContext flowContext = new FlowContext(flowId, flowInstanceId);
        flowContext.setFlow(flow);
        //set flowContext
        engineContext.setFlowContext(flowContext);
        //set input express evaluator
        engineContext.setInputParamEvaluators(flowRTObject.getInputParamEvaluators());
        //add flow
        engineContext.addFlow(flow);
        return engineContext;
    }

    public interface FlowRTObjectGetter {
        FlowRTObject get(String flowId);
    }

}